import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
import pickle
import os

# 路径配置
DATA_PATH = 'data/diabetes.csv'
MODEL_PATH = 'model/model.pkl'

# 读取数据
df = pd.read_csv(DATA_PATH)

# 特征处理
# 1. 将 0 视为缺失值的列填充为均值（按实际列名）
zero_as_nan_cols = ['plas', 'pres', 'skin', 'insu', 'mass']
for col in zero_as_nan_cols:
    df[col] = df[col].replace(0, pd.NA)
    mean_val = df[col].mean()
    df[col] = df[col].fillna(mean_val)

# 2. 目标编码
df['class'] = df['class'].map({'tested_negative': 0, 'tested_positive': 1})

# 3. 标准化所有特征
X = df.drop('class', axis=1)            # 删除目标列
scaler = StandardScaler()               # 标准化
X_scaled = scaler.fit_transform(X)      # 标准化
y = df['class']                         # 目标列  

# 训练模型
clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_scaled, y)

# 保存模型
os.makedirs('model', exist_ok=True)
with open(MODEL_PATH, 'wb') as f:
    pickle.dump(clf, f)

print('模型已训练并保存到 model/model.pkl')